import os
from os.path import join, isdir, isfile
import trimesh
import open3d as o3d
import numpy as np
from mesh import save_mesh


def simplify_mesh(o_trimesh):
    # TODO: 目前要求a和b必须是缩放的关系
    a_o3d = o3d.geometry.TriangleMesh(vertices=o3d.utility.Vector3dVector(o_trimesh.vertices), triangles=o3d.utility.Vector3iVector(o_trimesh.faces))
    if o_trimesh.faces.shape[0] > 2000:
        simple_a_o3d = a_o3d.simplify_quadric_decimation(target_number_of_triangles=o_trimesh.faces.shape[0] // 20)
    else:
        simple_a_o3d = a_o3d
    
    simple_trimesh = trimesh.Trimesh(vertices=np.float32(simple_a_o3d.vertices), faces=np.int32(simple_a_o3d.triangles))
    return simple_trimesh


def simplify_dataset_object_meshes():
    obj_model_root = "/share/datasets/HHO_object_dataset_final"
    save_root = "/data2/datasets/HHO_object_dataset_final_simplified"
    for category in os.listdir(obj_model_root):
        category_dir = join(obj_model_root, category)
        if not isdir(category_dir):
            continue
        for obj_name in os.listdir(category_dir):
            obj_dir = join(category_dir, obj_name)
            if not isdir(obj_dir):
                continue
            for file_name in os.listdir(obj_dir):
                if file_name.endswith("_m.obj"):
                    mesh = trimesh.load(join(obj_dir, file_name))
                    s_mesh = simplify_mesh(mesh)
                    v = s_mesh.vertices  # (N, 3)
                    if np.linalg.norm(v.max(axis=0)-v.min(axis=0), ord=2) > 2:
                        s_mesh.vertices *= 0.1
                    print(file_name, v.min(axis=0), v.max(axis=0), (v.min(axis=0)+v.max(axis=0))/2)
                    os.makedirs(join(save_root, category, obj_name), exist_ok=True)
                    save_mesh(s_mesh, join(save_root, category, obj_name, file_name))


if __name__ == "__main__":
    simplify_dataset_object_meshes()
